{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c3e7b0bf-5551-4256-a245-b6ca86e61949",
   "metadata": {},
   "source": [
    "# CUDA 测试代码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7229a0e0-296d-45fc-9683-faaea2e1bbf5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://mirrors.ustc.edu.cn/pypi/simple\n",
      "Requirement already satisfied: numba in /opt/app-root/lib/python3.9/site-packages (0.60.0)\n",
      "Requirement already satisfied: llvmlite<0.44,>=0.43.0dev0 in /opt/app-root/lib/python3.9/site-packages (from numba) (0.43.0)\n",
      "Requirement already satisfied: numpy<2.1,>=1.22 in /opt/app-root/lib/python3.9/site-packages (from numba) (1.24.4)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "CUDA is available:  False\n",
      "Modified array (using CUDA when available):  [ 10.  11.  12.  13.  14.  15.  16.  17.  18.  19.  20.  21.  22.  23.\n",
      "  24.  25.  26.  27.  28.  29.  30.  31.  32.  33.  34.  35.  36.  37.\n",
      "  38.  39.  40.  41.  42.  43.  44.  45.  46.  47.  48.  49.  50.  51.\n",
      "  52.  53.  54.  55.  56.  57.  58.  59.  60.  61.  62.  63.  64.  65.\n",
      "  66.  67.  68.  69.  70.  71.  72.  73.  74.  75.  76.  77.  78.  79.\n",
      "  80.  81.  82.  83.  84.  85.  86.  87.  88.  89.  90.  91.  92.  93.\n",
      "  94.  95.  96.  97.  98.  99. 100. 101. 102. 103. 104. 105. 106. 107.\n",
      " 108. 109.]\n",
      "Expected result:  [ 10.  11.  12.  13.  14.  15.  16.  17.  18.  19.  20.  21.  22.  23.\n",
      "  24.  25.  26.  27.  28.  29.  30.  31.  32.  33.  34.  35.  36.  37.\n",
      "  38.  39.  40.  41.  42.  43.  44.  45.  46.  47.  48.  49.  50.  51.\n",
      "  52.  53.  54.  55.  56.  57.  58.  59.  60.  61.  62.  63.  64.  65.\n",
      "  66.  67.  68.  69.  70.  71.  72.  73.  74.  75.  76.  77.  78.  79.\n",
      "  80.  81.  82.  83.  84.  85.  86.  87.  88.  89.  90.  91.  92.  93.\n",
      "  94.  95.  96.  97.  98.  99. 100. 101. 102. 103. 104. 105. 106. 107.\n",
      " 108. 109.]\n"
     ]
    }
   ],
   "source": [
    "!pip install numba -i https://mirrors.ustc.edu.cn/pypi/simple\n",
    "from numba import cuda, njit\n",
    "import numpy as np\n",
    "\n",
    "# 检查是否有可用的 GPU\n",
    "cuda_available = cuda.is_available()\n",
    "print(\"CUDA is available: \", cuda_available)\n",
    "\n",
    "\n",
    "def add_kernel_cpu(io_array, value):\n",
    "    for idx in range(io_array.size):\n",
    "        io_array[idx] += value\n",
    "\n",
    "\n",
    "# 定义一个简单的 CUDA 核函数（仅当CUDA可用时）\n",
    "if cuda_available:\n",
    "    @cuda.jit\n",
    "    def add_kernel_gpu(io_array, value):\n",
    "        idx = cuda.grid(1)\n",
    "        if idx < io_array.size:\n",
    "            io_array[idx] += value\n",
    "\n",
    "# 创建一个numpy数组\n",
    "data = np.arange(100, dtype=np.float32)\n",
    "data_original = data.copy()  # 保留原始数据以备后用\n",
    "\n",
    "if cuda_available:\n",
    "    # 分配GPU内存并将数据复制到GPU\n",
    "    data_gpu = cuda.to_device(data)\n",
    "\n",
    "    # 设置线程块大小和网格大小\n",
    "    threadsperblock = 256\n",
    "    blockspergrid = (data.size + (threadsperblock - 1)) // threadsperblock\n",
    "\n",
    "    # 调用核函数\n",
    "    add_kernel_gpu[blockspergrid, threadsperblock](data_gpu, 10)\n",
    "\n",
    "    # 将修改后的数据从GPU复制回主机\n",
    "    data_gpu.copy_to_host(data)\n",
    "else:\n",
    "    # 使用CPU进行计算\n",
    "    add_kernel_cpu(data, 10)\n",
    "\n",
    "print(\"Modified array (using CUDA when available): \", data)\n",
    "\n",
    "# 如果你想要验证结果是否正确，可以对比使用NumPy直接加10的结果\n",
    "expected_result = data_original + 10\n",
    "print(\"Expected result: \", expected_result)\n",
    "assert np.allclose(data, expected_result), \"The results do not match the expected output.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f506a7f0-e47b-4acc-8d87-0863ce84e7ab",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.16",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
